#pragma once
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/common/common.h>
#include <iostream>
#include <vector>
#include <array>
#include <stack>
#include <thread>
#include <algorithm>
#include <numeric>
#include <cmath>
#include <memory>
#include <mutex>
#include <functional>
#include <future>

using PointT = pcl::PointXYZ;
using PointCloud = pcl::PointCloud<PointT>;

// 包围盒结构
struct BoundingBox {
    std::array<float, 3> center; // 中心坐标
    float size;                  // 边长

    bool contains(const PointT& point) const {
        for (int i = 0; i < 3; ++i) {
            if (std::abs(point.data[i] - center[i]) > size / 2.0f) {
                return false;
            }
        }
        return true;
    }

    bool intersects(const BoundingBox& other) const {
        for (int i = 0; i < 3; ++i) {
            if (std::abs(center[i] - other.center[i]) > (size + other.size) / 2.0f) {
                return false;
            }
        }
        return true;
    }
};

// 对象池模板
template <typename T>
class ObjectPool {
public:
    ObjectPool(size_t chunkSize = 1024) : chunkSize(chunkSize) {}

    T* allocate() {
        if (freeList.empty()) {
            allocateChunk();
        }
        T* obj = freeList.back();
        freeList.pop_back();
        return obj;
    }

    void deallocate(T* obj) {
        freeList.push_back(obj);
    }

private:
    std::vector<T*> freeList;
    size_t chunkSize;

    void allocateChunk() {
        for (size_t i = 0; i < chunkSize; ++i) {
            freeList.push_back(new T());
        }
    }
};

// 八叉树节点
template <typename T>
class OctreeNode {
public:
    static ObjectPool<OctreeNode<T>> pool; // 静态对象池

    std::array<OctreeNode<T>*, 8> children{nullptr};
    std::vector<int> pointIndices; // 存储点索引
    size_t pointCount = 0;
    BoundingBox bbox; // 节点的包围盒

    OctreeNode() = default;

    // 判断是否有子节点
    bool hasChildren() const {
        for (int i = 0; i < 8; ++i) {
            if (children[i] != nullptr) {
                return true;
            }
        }
        return false;
    }

    // 从对象池分配/释放节点
    static OctreeNode<T>* createNode() { return pool.allocate(); }
    static void destroyNode(OctreeNode<T>* node) { pool.deallocate(node); }
};

template <typename T>
ObjectPool<OctreeNode<T>> OctreeNode<T>::pool;

// 八叉树类
class Octree {
public:
    Octree(float resolution, size_t densityThreshold)
        : resolution(resolution), densityThreshold(densityThreshold) {}

    ~Octree() {
        destroyTree(root);
    }

    // 根据点云自动计算根节点的包围盒并构建八叉树
    void buildTree(const PointCloud::Ptr& cloud) {
        this->cloud = cloud;

        // 计算根节点的包围盒
        rootBBox = computeBoundingBox(cloud);
        root = OctreeNode<PointT>::createNode();
        root->bbox = rootBBox;

        // 构建八叉树
        std::vector<size_t> indices(cloud->size());
        std::iota(indices.begin(), indices.end(), 0);
        std::sort(indices.begin(), indices.end(), [&](size_t a, size_t b) {
            return computeMorton(cloud->points[a]) < computeMorton(cloud->points[b]);
        });

        for (auto idx : indices) {
            insert(root, idx, root->bbox, 0);
        }
    }

    // 获取八叉树的最大深度
    int getMaxDepth() const {
        return maxDepth;
    }

    // 范围查询
    std::vector<int> queryRange(const BoundingBox& range) const {
        std::vector<int> results;
        std::stack<OctreeNode<PointT>*> stack;
        stack.push(root);

        while (!stack.empty()) {
            OctreeNode<PointT>* node = stack.top();
            stack.pop();

            if (!node->bbox.intersects(range)) continue;

            for (const auto& idx : node->pointIndices) {
                if (range.contains(cloud->points[idx])) {
                    results.push_back(idx);
                }
            }

            for (int i = 0; i < 8; ++i) {
                if (node->children[i]) {
                    stack.push(node->children[i]);
                }
            }
        }
        return results;
    }

    // 遍历所有节点（支持深度限制）
    void traverse(const std::function<void(OctreeNode<PointT>*, int depth)>& callback, int maxDepth = std::numeric_limits<int>::max()) const {
        std::stack<std::pair<OctreeNode<PointT>*, int>> stack;
        stack.push({root, 0});

        while (!stack.empty()) {
            auto [node, depth] = stack.top();
            stack.pop();

            // 调用回调函数
            callback(node, depth);

            // 如果达到最大深度，跳过子节点
            if (depth >= maxDepth) continue;

            // 遍历子节点
            for (int i = 0; i < 8; ++i) {
                if (node->children[i]) {
                    stack.push({node->children[i], depth + 1});
                }
            }
        }
    }

    // 并行遍历所有节点
    void parallelTraverse(const std::function<void(OctreeNode<PointT>*, int depth)>& callback, int maxDepth = std::numeric_limits<int>::max()) const {
        std::vector<std::future<void>> futures;
        std::mutex mutex;

        // 遍历根节点的子节点
        for (int i = 0; i < 8; ++i) {
            if (root->children[i]) {
                futures.push_back(std::async(std::launch::async, [&, i]() {
                    std::stack<std::pair<OctreeNode<PointT>*, int>> stack;
                    stack.push({root->children[i], 1});

                    while (!stack.empty()) {
                        auto [node, depth] = stack.top();
                        stack.pop();

                        // 调用回调函数
                        {
                            std::lock_guard<std::mutex> lock(mutex);
                            callback(node, depth);
                        }

                        // 如果达到最大深度，跳过子节点
                        if (depth >= maxDepth) continue;

                        // 遍历子节点
                        for (int j = 0; j < 8; ++j) {
                            if (node->children[j]) {
                                stack.push({node->children[j], depth + 1});
                            }
                        }
                    }
                }));
            }
        }

        // 等待所有任务完成
        for (auto& future : futures) {
            future.wait();
        }
    }

private:
    OctreeNode<PointT>* root = nullptr;
    BoundingBox rootBBox;
    float resolution; // 最小节点大小（分辨率）
    size_t densityThreshold;
    int maxDepth = 0; // 最大深度
    PointCloud::Ptr cloud;

    // 计算点云的包围盒
    BoundingBox computeBoundingBox(const PointCloud::Ptr& cloud) {
        if (cloud->empty()) {
            throw std::runtime_error("Point cloud is empty!");
        }

        // 初始化最小和最大坐标
        std::array<float, 3> minPt = {cloud->points[0].x, cloud->points[0].y, cloud->points[0].z};
        std::array<float, 3> maxPt = minPt;

        // 遍历点云，找到最小和最大坐标
        for (const auto& point : cloud->points) {
            for (int i = 0; i < 3; ++i) {
                if (point.data[i] < minPt[i]) minPt[i] = point.data[i];
                if (point.data[i] > maxPt[i]) maxPt[i] = point.data[i];
            }
        }

        // 计算中心点和边长
        BoundingBox bbox;
        for (int i = 0; i < 3; ++i) {
            bbox.center[i] = (minPt[i] + maxPt[i]) / 2.0f;
        }
        bbox.size = std::max({maxPt[0] - minPt[0], maxPt[1] - minPt[1], maxPt[2] - minPt[2]});

        return bbox;
    }

    void insert(OctreeNode<PointT>* node, int pointIdx, const BoundingBox& bbox, int depth) {
        // 更新最大深度
        if (depth > maxDepth) {
            maxDepth = depth;
        }

        // 如果节点大小小于分辨率或点数量小于阈值，停止分割
        if (bbox.size <= resolution || node->pointCount < densityThreshold) {
            node->pointIndices.push_back(pointIdx);
            node->pointCount++;
            return;
        }

        // 如果节点没有子节点，分割节点
        if (!node->hasChildren()) {
            splitNode(node, bbox);
        }

        // 计算点所在的子节点
        int octant = calculateOctant(cloud->points[pointIdx], bbox.center);
        insert(node->children[octant], pointIdx, node->children[octant]->bbox, depth + 1);
    }

    void splitNode(OctreeNode<PointT>* node, const BoundingBox& parentBBox) {
        for (int i = 0; i < 8; ++i) {
            node->children[i] = OctreeNode<PointT>::createNode();
            node->children[i]->bbox = getChildBBox(i, parentBBox);
        }
    }

    int calculateOctant(const PointT& point, const std::array<float, 3>& center) const {
        int octant = 0;
        for (int i = 0; i < 3; ++i) {
            if (point.data[i] >= center[i]) {
                octant |= (1 << i);
            }
        }
        return octant;
    }

    BoundingBox getChildBBox(int octant, const BoundingBox& parentBBox) const {
        BoundingBox childBBox = parentBBox;
        float newSize = parentBBox.size / 2.0f;
        for (int i = 0; i < 3; ++i) {
            childBBox.center[i] += (octant & (1 << i)) ? newSize / 2.0f : -newSize / 2.0f;
        }
        childBBox.size = newSize;
        return childBBox;
    }

    uint64_t computeMorton(const PointT& point) const {
        auto expandBits = [](uint32_t v) -> uint64_t {
            v = (v | (v << 16)) & 0x030000FF;
            v = (v | (v << 8)) & 0x0300F00F;
            v = (v | (v << 4)) & 0x030C30C3;
            v = (v | (v << 2)) & 0x09249249;
            return v;
        };

        uint32_t ix = static_cast<uint32_t>((point.x - rootBBox.center[0] + rootBBox.size / 2) / rootBBox.size * (1 << 21));
        uint32_t iy = static_cast<uint32_t>((point.y - rootBBox.center[1] + rootBBox.size / 2) / rootBBox.size * (1 << 21));
        uint32_t iz = static_cast<uint32_t>((point.z - rootBBox.center[2] + rootBBox.size / 2) / rootBBox.size * (1 << 21));
        return (expandBits(ix) << 2) | (expandBits(iy) << 1) | expandBits(iz);
    }

    void destroyTree(OctreeNode<PointT>* node) {
        if (!node) return;
        for (int i = 0; i < 8; ++i) {
            destroyTree(node->children[i]);
        }
        OctreeNode<PointT>::destroyNode(node);
    }
};

// 测试代码
// int test() {
//     PointCloud::Ptr cloud(new PointCloud);
//     for (int i = 0; i < 1000000; ++i) {
//         cloud->push_back(PointT(static_cast<float>(rand()) / RAND_MAX, static_cast<float>(rand()) / RAND_MAX, static_cast<float>(rand()) / RAND_MAX));
//     }

//     Octree octree(0.25f, 10); // 分辨率为 0.01，密度阈值为 10
//     octree.buildTree(cloud);

//     // 获取最大深度
//     std::cout << "Max depth: " << octree.getMaxDepth() << std::endl;

//     // 单线程遍历所有节点
//     octree.traverse([](OctreeNode<PointT>* node, int depth) {
//         std::cout << "Node center: (" << node->bbox.center[0] << ", " << node->bbox.center[1] << ", " << node->bbox.center[2] << "), "
//                   << "Size: " << node->bbox.size << ", "
//                   << "Depth: " << depth << ", "
//                   << "Points: " << node->pointIndices.size() << std::endl;
//     });

//     return 0;
// }
